import json
import os
from .onnx_model import KittenTTS_1_Onnx


class KittenTTS:
    """Main KittenTTS class for text-to-speech synthesis."""
    
    def __init__(self, model_name="KittenML/kitten-tts-nano-0.2", cache_dir=None, language="en-us", espeak_path=None):
        """Initialize KittenTTS with a model from Hugging Face.
        
        Args:
            model_name: Hugging Face repository ID or model name
            cache_dir: Directory to cache downloaded files
            language: Language code for phonemization (e.g. "en-us" for English, "zh" for Chinese)
            espeak_path: Path to the espeak-ng executable (optional)
        """
        # Handle different model name formats
        if "/" not in model_name:
            # If just model name provided, assume it's from KittenML
            repo_id = f"KittenML/{model_name}"
        else:
            repo_id = model_name
            
        self.model = download_from_huggingface(repo_id=repo_id, cache_dir=cache_dir, language=language, espeak_path=espeak_path)
    
    def generate(self, text, voice="expr-voice-5-m", speed=1.0):
        """Generate audio from text.
        
        Args:
            text: Input text to synthesize
            voice: Voice to use for synthesis
            speed: Speech speed (1.0 = normal)
            
        Returns:
            Audio data as numpy array
        """
        return self.model.generate(text, voice=voice, speed=speed)
    
    def generate_to_file(self, text, output_path, voice="expr-voice-5-m", speed=1.0, sample_rate=24000):
        """Generate audio from text and save to file.
        
        Args:
            text: Input text to synthesize
            output_path: Path to save the audio file
            voice: Voice to use for synthesis
            speed: Speech speed (1.0 = normal)
            sample_rate: Audio sample rate
        """
        return self.model.generate_to_file(text, output_path, voice=voice, speed=speed, sample_rate=sample_rate)
    
    @property
    def available_voices(self):
        """Get list of available voices."""
        return self.model.available_voices


def download_from_huggingface(repo_id="KittenML/kitten-tts-nano-0.1", cache_dir=None, language="en-us", espeak_path=None):
    """Load model files from local directory.
    
    Args:
        repo_id: Model repository ID (now used to determine local path)
        cache_dir: Not used anymore, kept for backward compatibility
        language: Language code for phonemization
        espeak_path: Path to the espeak-ng executable (optional)
        
    Returns:
        KittenTTS_1_Onnx: Instantiated model ready for use
    """
    # 从repo_id中提取模型目录名
    if "/" in repo_id:
        # 如果是完整的仓库ID格式，提取最后一部分作为模型目录名
        model_dir_name = repo_id.split("/")[-1]
        # 处理类似 "KittenML/kitten-tts-nano-0.2" 的格式
        if repo_id.split("/")[0] == "KittenML":
            # 保持原有的KittenML目录结构
            local_model_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "KittenML", model_dir_name)
        else:
            # 对于其他仓库，直接使用模型目录名
            local_model_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), model_dir_name)
    else:
        # 如果只有模型名称，使用它作为模型目录名
        local_model_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), repo_id)
    
    # Path to config file
    config_path = os.path.join(local_model_dir, "config.json")
    
    # Check if config file exists
    if not os.path.exists(config_path):
        raise FileNotFoundError(f"Config file not found at {config_path}")
    
    # Load config
    with open(config_path, 'r') as f:
        config = json.load(f)

    if config.get("type") != "ONNX1":
        raise ValueError("Unsupported model type.")

    # Get model and voices files from config
    model_path = os.path.join(local_model_dir, config["model_file"])
    voices_path = os.path.join(local_model_dir, config["voices"])
    
    # Check if model and voices files exist
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Model file not found at {model_path}")
    if not os.path.exists(voices_path):
        raise FileNotFoundError(f"Voices file not found at {voices_path}")
    
    # Instantiate and return model
    model = KittenTTS_1_Onnx(model_path=model_path, voices_path=voices_path, language=language, espeak_path=espeak_path)
    
    return model


def get_model(repo_id="KittenML/kitten-tts-nano-0.1", cache_dir=None, language="en-us", espeak_path=None):
    """Get a KittenTTS model (legacy function for backward compatibility)."""
    return KittenTTS(repo_id, cache_dir, language, espeak_path)
